
// Copyright (C) Recursion Company. All rights reserved.

#include "stdafx.h"
#include "vol_user_app_info.h"
#include "vol_des.h"

typedef struct rc4_key
{
    unsigned char state [256];
    unsigned char x;
    unsigned char y;
} rc4_key;

static void swap_byte (unsigned char* a, unsigned char* b)
{
   unsigned char swapByte;

   swapByte = *a;
   *a = *b;
   *b = swapByte;
}

static void prepare_key (unsigned char* key_data_ptr, int key_data_len, rc4_key* key)
{
    unsigned char index1;
    unsigned char index2;
    unsigned char* state;
    int counter;

    state = &key->state [0];
    for (counter = 0; counter < 256; counter++)
        state [counter] = (unsigned char)counter;

    key->x = 0;
    key->y = 0;

    if (key_data_len > 0)
    {
        index1 = 0;
        index2 = 0;
        for (counter = 0; counter < 256; counter++)
        {
            index2 = (key_data_ptr [index1] + state [counter] + index2) % 256;
            swap_byte (&state [counter], &state [index2]);

            index1 = (index1 + 1) % key_data_len;
        }
    }
}

static void rc4 (unsigned char* buffer_ptr, int buffer_len, rc4_key* key)
{
    unsigned char x;
    unsigned char y;
    unsigned char* state;
    unsigned char xorIndex;
    int counter;

    x = key->x;
    y = key->y;

    state = &key->state [0];
    for (counter = 0; counter < buffer_len; counter++)
    {
        x = (x + 1) % 256;
        y = (state [x] + y) % 256;
        swap_byte (&state [x], &state [y]);

        xorIndex = (state [x] + state [y]) % 256;
        buffer_ptr [counter] ^= state [xorIndex];
    }

    key->x = x;     
    key->y = y;
}

CVolMem& VolRunEncrypt (const BYTE* pbData, INT nDataSize, const TCHAR* szPassword, BOOL blIsEncrypt, INT nMode, CVolMem& memResult)
{
    memResult.Empty ();

    if (nDataSize <= 0)
        return memResult;

    CVolMem memBuf;
    CHAR* szAnsiPassword = (CHAR*)GetMbsText (szPassword, memBuf, NULL);

    if (nMode == 1)
    {
        if (blIsEncrypt)
        {
            INT n = (nDataSize + sizeof (INT)) % 8;
            if (n > 0)
                n = 8 - n;
            INT nOutSize = sizeof (INT) + nDataSize + n;

            LPBYTE pOut = (LPBYTE)memResult.Alloc (nOutSize, TRUE);
            *(LPINT)pOut = nDataSize;
            COPY_MEM (pOut + sizeof (INT), pbData, nDataSize);
            
            vol_des_encrypt (pOut, nOutSize, szAnsiPassword);
        }
        else if ((nDataSize % 8) == 0)
        {
            LPBYTE pOut = (LPBYTE)memResult.Alloc (nDataSize);

            COPY_MEM (pOut, pbData, nDataSize);
            vol_des_decrypt (pOut, nDataSize, szAnsiPassword);

            INT nOutSize = *(LPINT)pOut;
            if (nOutSize <= 0 || nDataSize >= nOutSize + 12 || nOutSize > nDataSize - 4)
            {
                memResult.Free ();
            }
            else
            {
                memResult.Remove (0, sizeof (INT));
                memResult.Realloc (nOutSize);
            }
        }
    }
    else
    {
        rc4_key key;
        prepare_key ((unsigned char*)szAnsiPassword, strlen (szAnsiPassword), &key);

        LPBYTE pOut = (LPBYTE)memResult.Alloc (nDataSize);
        COPY_MEM (pOut, pbData, nDataSize);

        rc4 (pOut, nDataSize, &key);
    }

    return memResult;
}
